{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'17485': 0, '2456': 1, '2730': 2, '3001': 3, '3002': 4, '3003': 5, '3004': 6, '3006': 7, '3007': 8, '3009': 9, '3010': 10, '30414': 11, '3622': 12, '45176': 13, '4600': 14, '55981': 15, '56145': 16, '56891': 17, '57520': 18, '6111': 19, '70695': 20, '87079': 21}\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[17], line 105\u001b[0m\n\u001b[0;32m    103\u001b[0m \u001b[38;5;66;03m# 保存模型\u001b[39;00m\n\u001b[0;32m    104\u001b[0m model_save_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m../model\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m--> 105\u001b[0m \u001b[43mtrain_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcriterion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtest_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mEPOCHS\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mDEVICE\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel_save_path\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    106\u001b[0m torch\u001b[38;5;241m.\u001b[39msave(model\u001b[38;5;241m.\u001b[39mstate_dict(), model_save_path)\n",
      "Cell \u001b[1;32mIn[17], line 79\u001b[0m, in \u001b[0;36mtrain_model\u001b[1;34m(model, criterion, optimizer, train_loader, test_loader, epochs, device, model_save_path)\u001b[0m\n\u001b[0;32m     77\u001b[0m     loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[0;32m     78\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m---> 79\u001b[0m     running_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitem\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     81\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mEpoch \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepoch\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepochs\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Training Loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrunning_loss\u001b[38;5;241m/\u001b[39m\u001b[38;5;28mlen\u001b[39m(train_loader)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     83\u001b[0m \u001b[38;5;66;03m# 测试模型\u001b[39;00m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# 模型训练\n",
    "import os\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision import datasets, models\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import random_split\n",
    "\n",
    "# 设置全局参数\n",
    "modellr = 1e-4\n",
    "BATCH_SIZE = 64\n",
    "EPOCHS = 50\n",
    "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# 数据预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.Grayscale(num_output_channels=1),  # 转换为单通道灰度图像\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5])\n",
    "\n",
    "])\n",
    "\n",
    "# 读取样本和标签\n",
    "data_set = datasets.ImageFolder('../Dataset/photos', transform)\n",
    "print(data_set.class_to_idx)\n",
    "\n",
    "# 随机划分数据集\n",
    "train_ratio = 0.8\n",
    "train_size = int(train_ratio * len(data_set))\n",
    "test_size = len(data_set) - train_size\n",
    "train_set, test_set = random_split(data_set, [train_size, test_size])\n",
    "\n",
    "# 导入数据\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False)\n",
    "\n",
    "# 实例化\n",
    "model = models.resnet50(pretrained=True)\n",
    "# 修改第一层以接受单通道图像\n",
    "model.conv1 = nn.Conv2d(1, model.conv1.out_channels, kernel_size=model.conv1.kernel_size, \n",
    "                        stride=model.conv1.stride, padding=model.conv1.padding, bias=False)\n",
    "\n",
    "# for param in model.parameters():\n",
    "#     param.requires_grad = False\n",
    "# # 解冻模型的最后两层\n",
    "# for name, param in model.named_parameters():\n",
    "#     if \"layer4\" in name or \"fc\" in name:\n",
    "#         param.requires_grad = True\n",
    "\n",
    "num_ftrs = model.fc.in_features\n",
    "num_classes = len(data_set.classes)\n",
    "model.fc = nn.Linear(num_ftrs, num_classes)\n",
    "model = model.to(DEVICE)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=modellr)  # 选择简单暴力的Adam优化器，学习率调低\n",
    "\n",
    "def adjust_learning_rate(optimizer, epoch, initial_lr):\n",
    "    \"\"\"Sets the learning rate to the initial LR decayed by 10 every 30 epochs\"\"\"\n",
    "    lr = initial_lr * (0.1 ** (epoch // 30))\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "\n",
    "\n",
    "def train_model(model, criterion, optimizer, train_loader, test_loader, epochs, device, model_save_path):\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        running_loss = 0.0\n",
    "        for inputs, labels in train_loader:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            running_loss += loss.item()\n",
    "\n",
    "        print(f\"Epoch {epoch+1}/{epochs}, Training Loss: {running_loss/len(train_loader)}\")\n",
    "\n",
    "        # 测试模型\n",
    "        model.eval()\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        with torch.no_grad():\n",
    "            for inputs, labels in test_loader:\n",
    "                inputs, labels = inputs.to(device), labels.to(device)\n",
    "                outputs = model(inputs)\n",
    "                _, predicted = torch.max(outputs.data, 1)\n",
    "                total += labels.size(0)\n",
    "                correct += (predicted == labels).sum().item()\n",
    "\n",
    "        accuracy = 100 * correct / total\n",
    "        print(f\"Epoch {epoch+1}/{epochs}, Test Accuracy: {accuracy:.2f}%\")\n",
    "\n",
    "        # 保存模型的 state_dict\n",
    "        torch.save(model.state_dict(), os.path.join(model_save_path, f\"model_final3_epoch_{epoch+1}.pth\"))\n",
    "\n",
    "    print(\"Training and testing complete\")\n",
    "\n",
    "# 保存模型\n",
    "model_save_path = \"../model\"\n",
    "train_model(model, criterion, optimizer, train_loader, test_loader, EPOCHS, DEVICE, model_save_path)\n",
    "torch.save(model.state_dict(), model_save_path)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'17485': 0, '2456': 1, '2730': 2, '3001': 3, '3002': 4, '3003': 5, '3004': 6, '3006': 7, '3007': 8, '3009': 9, '3010': 10, '30414': 11, '3622': 12, '45176': 13, '4600': 14, '55981': 15, '56145': 16, '56891': 17, '57520': 18, '6111': 19, '70695': 20, '87079': 21}\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[13], line 116\u001b[0m\n\u001b[0;32m    114\u001b[0m \u001b[38;5;66;03m# 保存模型\u001b[39;00m\n\u001b[0;32m    115\u001b[0m model_save_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m../model\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m--> 116\u001b[0m \u001b[43mtrain_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcriterion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtest_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mEPOCHS\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mDEVICE\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel_save_path\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    117\u001b[0m torch\u001b[38;5;241m.\u001b[39msave(model\u001b[38;5;241m.\u001b[39mstate_dict(), model_save_path)\n",
      "Cell \u001b[1;32mIn[13], line 90\u001b[0m, in \u001b[0;36mtrain_model\u001b[1;34m(model, criterion, optimizer, train_loader, test_loader, epochs, device, model_save_path)\u001b[0m\n\u001b[0;32m     88\u001b[0m     loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[0;32m     89\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m---> 90\u001b[0m     running_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitem\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     92\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mEpoch \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepoch\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mepochs\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Training Loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrunning_loss\u001b[38;5;241m/\u001b[39m\u001b[38;5;28mlen\u001b[39m(train_loader)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     94\u001b[0m \u001b[38;5;66;03m# 测试模型\u001b[39;00m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# 模型迁移学习\n",
    "import os\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision import datasets, models\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import random_split\n",
    "\n",
    "# 设置全局参数\n",
    "modellr = 1e-4\n",
    "BATCH_SIZE = 64\n",
    "EPOCHS = 20\n",
    "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "# 模型路径\n",
    "model_path = '../model/model_epoch_3.pth' \n",
    "image_dir = '../Dataset/test/7'\n",
    "output_dir = '../Dataset/test_result' \n",
    "\n",
    "# 数据预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.Grayscale(num_output_channels=1),  # 转换为单通道灰度图像\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5])\n",
    "\n",
    "])\n",
    "\n",
    "# 读取样本和标签\n",
    "data_set = datasets.ImageFolder('../Dataset/photos', transform)\n",
    "print(data_set.class_to_idx)\n",
    "\n",
    "# 随机划分数据集\n",
    "train_ratio = 0.8\n",
    "train_size = int(train_ratio * len(data_set))\n",
    "test_size = len(data_set) - train_size\n",
    "train_set, test_set = random_split(data_set, [train_size, test_size])\n",
    "\n",
    "# 导入数据\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False)\n",
    "\n",
    "# 创建模型实例\n",
    "model = models.resnet50(pretrained=False)\n",
    "# 修改第一层以接受单通道图像\n",
    "model.conv1 = nn.Conv2d(1, model.conv1.out_channels, kernel_size=model.conv1.kernel_size, \n",
    "                        stride=model.conv1.stride, padding=model.conv1.padding, bias=False)\n",
    "# 冻结前面的层\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "# 解冻模型的最后两层\n",
    "for name, param in model.named_parameters():\n",
    "    if \"layer4\" in name or \"fc\" in name:\n",
    "        param.requires_grad = True\n",
    "\n",
    "num_ftrs = model.fc.in_features\n",
    "num_classes = len(data_set.classes)\n",
    "model.fc = nn.Linear(num_ftrs, num_classes)\n",
    "\n",
    "# 删除最后一层权重后,加载模型的 state_dict\n",
    "state_dict = torch.load(model_path)\n",
    "state_dict.pop('fc.weight', None)\n",
    "state_dict.pop('fc.bias', None)\n",
    "model.load_state_dict(state_dict, strict=False)\n",
    "\n",
    "# 将模型转移到cuda上\n",
    "model = model.to(DEVICE)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=modellr)  # 选择简单暴力的Adam优化器，学习率调低\n",
    "\n",
    "def adjust_learning_rate(optimizer, epoch, initial_lr):\n",
    "    \"\"\"Sets the learning rate to the initial LR decayed by 10 every 30 epochs\"\"\"\n",
    "    lr = initial_lr * (0.1 ** (epoch // 30))\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "\n",
    "\n",
    "def train_model(model, criterion, optimizer, train_loader, test_loader, epochs, device, model_save_path):\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        running_loss = 0.0\n",
    "        for inputs, labels in train_loader:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            running_loss += loss.item()\n",
    "\n",
    "        print(f\"Epoch {epoch+1}/{epochs}, Training Loss: {running_loss/len(train_loader)}\")\n",
    "\n",
    "        # 测试模型\n",
    "        model.eval()\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        with torch.no_grad():\n",
    "            for inputs, labels in test_loader:\n",
    "                inputs, labels = inputs.to(device), labels.to(device)\n",
    "                outputs = model(inputs)\n",
    "                _, predicted = torch.max(outputs.data, 1)\n",
    "                total += labels.size(0)\n",
    "                correct += (predicted == labels).sum().item()\n",
    "\n",
    "        accuracy = 100 * correct / total\n",
    "        print(f\"Epoch {epoch+1}/{epochs}, Test Accuracy: {accuracy:.2f}%\")\n",
    "\n",
    "        # 保存模型的 state_dict\n",
    "        torch.save(model.state_dict(), os.path.join(model_save_path, f\"model_final2_epoch_{epoch+1}.pth\"))\n",
    "\n",
    "    print(\"Training and testing complete\")\n",
    "\n",
    "# 保存模型\n",
    "model_save_path = \"../model\"\n",
    "train_model(model, criterion, optimizer, train_loader, test_loader, EPOCHS, DEVICE, model_save_path)\n",
    "torch.save(model.state_dict(), model_save_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型预测使用\n",
    "**注意：输入彩色图像，下面程序自动转换为灰度图，该模型采用灰度图进行训练和预测，**\n",
    "将`image_dir`,`output_dir`都设置为合理的目录\n",
    "将`model_path`指定为模型的.pth文件(保存了模型的state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All images processed and saved in the output directory.\n"
     ]
    }
   ],
   "source": [
    "# 模型预测效果\n",
    "from PIL import Image, ImageDraw\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import transforms, models\n",
    "\n",
    "# 模型路径\n",
    "model_path = '../model/model_final2_epoch_1.pth' \n",
    "image_dir = '../Dataset/test/'\n",
    "output_dir = '../Dataset/test_result' \n",
    "num_classes = 22\n",
    "\n",
    "# 创建模型实例\n",
    "model = models.resnet50(pretrained=False)    # 这里好像错了，不应该加载预先训练权重的\n",
    "# 修改第一层以接受单通道图像\n",
    "model.conv1 = nn.Conv2d(1, model.conv1.out_channels, kernel_size=model.conv1.kernel_size, \n",
    "                        stride=model.conv1.stride, padding=model.conv1.padding, bias=False)\n",
    "\n",
    "num_ftrs = model.fc.in_features\n",
    "model.fc = torch.nn.Linear(num_ftrs, num_classes)  # <num_classes> 需要替换为类别数量\n",
    "\n",
    "# 加载模型的 state_dict\n",
    "model.load_state_dict(torch.load(model_path))\n",
    "\n",
    "model.eval()\n",
    "\n",
    "# 预处理函数\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.Grayscale(num_output_channels=1),  # 转换为单通道灰度图像\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5])\n",
    "])\n",
    "\n",
    "# 加载并预处理图片\n",
    "def load_and_preprocess_image(image_path):\n",
    "    image = Image.open(image_path)\n",
    "    image = transform(image).unsqueeze(0)\n",
    "    return image\n",
    "\n",
    "# 预测函数\n",
    "def predict(model, image_tensor):\n",
    "    outputs = model(image_tensor)\n",
    "    _, predicted = torch.max(outputs.data, 1)\n",
    "    return predicted.item()\n",
    "\n",
    "# 对文件夹中的每张图像进行预测，并在图像上写入预测结果\n",
    "for filename in os.listdir(image_dir):\n",
    "    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):\n",
    "        image_path = os.path.join(image_dir, filename)\n",
    "        image_tensor = load_and_preprocess_image(image_path)\n",
    "        predicted_class_idx = predict(model, image_tensor)\n",
    "\n",
    "        # 在图像上写入预测结果\n",
    "        output_image = Image.open(image_path)\n",
    "        draw = ImageDraw.Draw(output_image)\n",
    "        draw.text((10, 30), f'Pre: {predicted_class_idx}', fill=(255, 255, 255))\n",
    "\n",
    "        # 保存修改后的图像\n",
    "        output_path = os.path.join(output_dir, filename)\n",
    "        output_image.save(output_path)\n",
    "\n",
    "print(\"All images processed and saved in the output directory.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
